#Power as Function of Republican Treatment Effect
library(ggplot2)
library(lmtest)
library(sandwich)

rm(list = ls())
set.seed(123)
Num_Sims <- 5000
Num_Steps <- 30
Num_Respondents <-3000;
Pr_Treatment <- 0.5

pref_param <- runif(Num_Respondents, min = 0, max = 1)

#Assuming Approximately 0.35 Republicans
Pr_Rep_Baseline <- pnorm(pref_param, mean = 0.6441, sd = 0.1)

treatment_effect_size <- matrix(,nrow = 1, ncol = Num_Steps)
alphap1_power_Rep <-matrix(,nrow = 1, ncol = Num_Steps)
alphap1_power_App <-matrix(,nrow = 1, ncol = Num_Steps)
alphap1_power_App_Cond_Rep <-matrix(,nrow = 1, ncol = Num_Steps)
alphap1_power_App_Cond_NonRep <-matrix(,nrow = 1, ncol = Num_Steps)


tau_rep_steps <- -1*c(1:Num_Steps)/(10*Num_Steps)
for(j in 1:Num_Steps){
  Pr_Rep_Treatment <- pnorm(pref_param, mean = 0.6441+tau_rep_steps[j], sd = 0.1)
  treatment_effect_size[,j]<-mean(Pr_Rep_Treatment)-mean(Pr_Rep_Baseline)
  
  #Assuming Approximately Baseline 0.20 Non-Republican Approval Rate and 0.85 Republican Approval Rate;
  Pr_App_Baseline <- pnorm(pref_param, mean = 0.57, sd = 0.189)
  sum(Pr_App_Baseline*Pr_Rep_Baseline)/sum(Pr_Rep_Baseline)
  sum(Pr_App_Baseline*(1-Pr_Rep_Baseline))/sum(1-Pr_Rep_Baseline)
  
  #Assume No Effect on Probability Approval
  Pr_App_Treatment <- Pr_App_Baseline
  
  ATE_Rep = matrix(,nrow = 1, ncol = Num_Sims)
  SE_ATE_Rep = matrix(,nrow = 1, ncol = Num_Sims)
  
  ATE_App = matrix(,nrow = 1, ncol = Num_Sims)
  SE_ATE_App = matrix(,nrow = 1, ncol = Num_Sims)
  
  ATE_App_Cond_Rep = matrix(,nrow = 1, ncol = Num_Sims)
  SE_ATE_App_Cond_Rep = matrix(,nrow = 1, ncol = Num_Sims)
  
  ATE_App_Cond_NonRep <- matrix(, nrow = 1, ncol = Num_Sims)
  SE_ATE_App_Cond_NonRep <- matrix(, nrow = 1, ncol = Num_Sims)
  
  for (k  in 1:Num_Sims){
    
    App_Control <- rbinom(Num_Respondents, 1, Pr_App_Baseline)
    Rep_Control <- rbinom(Num_Respondents, 1, Pr_Rep_Baseline)
    
    App_Treatment <- rbinom(Num_Respondents, 1, Pr_App_Treatment)
    Rep_Treatment <- rbinom(Num_Respondents, 1, Pr_Rep_Treatment)
    
    Treatment_Indicator <-rbinom(Num_Respondents, 1, Pr_Treatment)
    
    Obs_Rep_Vector <- Rep_Treatment*(Treatment_Indicator) + (Rep_Control)*(1-Treatment_Indicator)
    Obs_App_Vector <- App_Treatment*(Treatment_Indicator) + (App_Control)*(1-Treatment_Indicator)
    
    Obs_App_Cond_Rep_Vector <- Obs_App_Vector[Obs_Rep_Vector == 1]
    Treatment_Indicator_Cond_Rep <- Treatment_Indicator[Obs_Rep_Vector == 1]
    
    Obs_App_Cond_NonRep_Vector <- Obs_App_Vector[Obs_Rep_Vector == 0]
    Treatment_Indicator_Cond_NonRep <- Treatment_Indicator[Obs_Rep_Vector == 0]
    
    m_rep_est_te <- lm(Obs_Rep_Vector ~ Treatment_Indicator)
    se_rep_est_te <- coeftest(m_rep_est_te, vcov = vcovHC(m_rep_est_te, type = "HC1"))
    
    ATE_Rep[,k] <- se_rep_est_te[2,1]
    SE_ATE_Rep[,k] <- se_rep_est_te[2,2]
    
    m_app_est_te <- lm(Obs_App_Vector ~ Treatment_Indicator)
    se_app_est_te <- coeftest(m_app_est_te, vcov = vcovHC(m_app_est_te, type = "HC1"))
    
    ATE_App[,k] <- se_app_est_te[2,1]
    SE_ATE_App[,k] <- se_app_est_te[2,2]
    
    m_app_cond_rep_est_te <- lm(Obs_App_Cond_Rep_Vector ~ Treatment_Indicator_Cond_Rep)
    se_app_cond_rep_est_te <-  coeftest(m_app_cond_rep_est_te, vcov = vcovHC(m_app_cond_rep_est_te, type = "HC1"))
    
    ATE_App_Cond_Rep[,k] <- se_app_cond_rep_est_te[2,1]
    SE_ATE_App_Cond_Rep[,k] <- se_app_cond_rep_est_te[2,2]
    
    m_app_cond_nonrep_est_te <- lm(Obs_App_Cond_NonRep_Vector ~ Treatment_Indicator_Cond_NonRep)
    se_app_cond_nonrep_est_te <-  coeftest(m_app_cond_nonrep_est_te, vcov = vcovHC(m_app_cond_nonrep_est_te, type = "HC1"))
    
    ATE_App_Cond_NonRep[,k] <- se_app_cond_nonrep_est_te[2,1]
    SE_ATE_App_Cond_NonRep[,k] <- se_app_cond_nonrep_est_te[2,2]
    
  }
  
  tstat_Rep <- ATE_Rep/SE_ATE_Rep
  tstat_App <- ATE_App/SE_ATE_App
  tstat_App_Cond_Rep <- ATE_App_Cond_Rep/SE_ATE_App_Cond_Rep
  tstat_App_Cond_NonRep <- ATE_App_Cond_NonRep/SE_ATE_App_Cond_NonRep
  
  #Level alpha = 0.1 two-sided tests
  alphap1_power_Rep[,j] <- mean(abs(tstat_Rep) > 1.644854)
  alphap1_power_App[,j] <- mean(abs(tstat_App) > 1.644854)
  alphap1_power_App_Cond_Rep[,j] <- mean(abs(tstat_App_Cond_Rep) > 1.644854)
  alphap1_power_App_Cond_NonRep[,j] <- mean(abs(tstat_App_Cond_NonRep) > 1.644854)
  
}

df_alphap1_power_Rep <- data.frame(cbind(t(treatment_effect_size), t(alphap1_power_Rep)))
plot_alphap1_power_Rep <- ggplot(df_alphap1_power_Rep, aes(x = X1, y = X2)) + geom_point(size=2) + ggtitle("Republican Self-Identification") + labs(x="True Effect Size", y = "Simulated Power of 0.1 Level Test") + theme_classic() + theme(text = element_text(size=12)) +  theme(plot.title = element_text(hjust = 0.5)) + geom_vline(xintercept=0.04) + ylim(0, 1)
ggsave("Output/plot_alphap1_power_Rep.pdf", plot_alphap1_power_Rep,
       width = 5, height = 4, dpi = 300, units = "in", device='pdf')

df_alphap1_power_App_Cond_Rep <- data.frame(cbind(t(treatment_effect_size), t(alphap1_power_App_Cond_Rep)))
plot_alphap1_power_App_Cond_Rep <- ggplot(df_alphap1_power_App_Cond_Rep, aes(x = X1, y = X2)) + geom_point(size=2) + ggtitle("Approval Among \n Self-Identified Republicans") + labs(x="True Effect Size", y = "Simulated Power of 0.1 Level Test") + theme_classic() + theme(text = element_text(size=12)) +  theme(plot.title = element_text(hjust = 0.5))  + geom_vline(xintercept=0.04) + ylim(0, 1)
ggsave("Output/plot_alphap1_power_App_Cond_Rep.pdf", plot_alphap1_power_App_Cond_Rep,
       width = 5, height = 4, dpi = 300, units = "in", device='pdf')

df_alphap1_power_App_Cond_NonRep <- data.frame(cbind(t(treatment_effect_size), t(alphap1_power_App_Cond_Rep)))
plot_alphap1_power_App_Cond_NonRep <- ggplot(df_alphap1_power_App_Cond_NonRep, aes(x = X1, y = X2)) + geom_point(size=2) + ggtitle("Approval Among \n Self-Identified Non-Republicans") + labs(x="True Effect Size", y = "Simulated Power of 0.1 Level Test") + theme_classic() + theme(text = element_text(size=12)) +  theme(plot.title = element_text(hjust = 0.5))  + geom_vline(xintercept=0.04) + ylim(0, 1)
ggsave("Output/plot_alphap1_power_App_Cond_NonRep.pdf", plot_alphap1_power_App_Cond_NonRep,
       width = 5, height = 4, dpi = 300, units = "in", device='pdf')